# -*- coding:utf-8 -*-

# @Time  : 2023/10/22 6:18 PM
# @Author: chenyong
# @Email : chenyong@lingxi.ai
# @File  : generate_bot_dialogue_by_deal.py

import tqdm
import utils
import random
import os
import pandas as pd
import uuid

def get_deal_data(path):
    datas = []
    for csv_path in path:
        if os.path.isdir(csv_path):
            for file in os.listdir(csv_path):
                if not file.endswith('.csv'):
                    continue
                print(csv_path + file)
                df = pd.read_csv(csv_path + file)
                faqDictList = df.to_dict('records')
                datas.extend(faqDictList)
        else:
            df = pd.read_csv(csv_path)
            faqDictList = df.to_dict('records')
            datas.extend(faqDictList)

    all_sessions = {}
    for index, faq in enumerate(datas):
        sessionId = str(faq['sessionId'])
        if sessionId not in all_sessions:
            all_sessions[sessionId] = []
            all_sessions[sessionId].append({"role": "用户", "value": "你好。"})

        if str(faq['过滤文本']) == '0':
            all_sessions[sessionId].append({"role": str(faq['角色']), "value": str(faq['原始文本']).split('#@')[-1].split('@@')[0]})
        else:
            all_sessions[sessionId].append({"role": str(faq['角色']), "value": str(faq['过滤文本'])})

    all_sessions_new = {}
    for session_id, session_value in all_sessions.items():
        new_session_value = []
        i = 0
        while i < len(session_value):
            j = i + 1
            while j < len(session_value):
                if session_value[j]['role'] == session_value[i]['role']:
                    session_value[i]['value'] += session_value[j]['value']
                elif session_value[i]['role'] == '用户':
                    new_session_value.append(session_value[i])
                    new_session_value.append(session_value[j])
                    break
                j += 1
            i = j + 1
        if len(new_session_value) > 14:
            all_sessions_new[session_id] = new_session_value
    all_data = []
    for _id, _conv in all_sessions_new.items():
        new_data_dict = {}
        new_data_dict['id'] = str(uuid.uuid1())
        new_data_dict['conversations'] = []
        i = 0
        while i < len(_conv):
            j = i + 1
            new_data_dict['conversations'].append({"from": "human", "value": _conv[i]['value']})
            new_data_dict['conversations'].append({"from": "gpt", "value": _conv[j]['value']})
            i = j + 1
        all_data.append(new_data_dict)

    print(len(all_data))
    return all_data


def write_json(path, data):
    utils.jdump(data, path)


if __name__ == '__main__':
    d = get_deal_data(["../deal_data/"])
    write_json("../data/deal_data.json", d)
